import streamlit as st
import psycopg2
import pandas as pd
from openai import OpenAI
from io import BytesIO

# ------------------------------
# 页面设置
# ------------------------------
st.set_page_config(
    page_title="pgsql_code_generator",
    page_icon="🔎",
    layout="wide",
    initial_sidebar_state="expanded",
)

# ------------------------------
# 初始化
# ------------------------------
client = OpenAI(
    api_key="xk-349774cce86e12913221634fbb69b071ec05e8c01d6d9d45",
    base_url="https://openai.zhixueyouke.cn/v1/"
)

if "conn" not in st.session_state:
    st.session_state.conn = None
if "table_structures" not in st.session_state:
    st.session_state.table_structures = None
if "selected_tables" not in st.session_state:
    st.session_state.selected_tables = []

# ------------------------------
# 工具函数
# ------------------------------

def connect_db(host, port, dbname, user, password):
    try:
        conn = psycopg2.connect(
            host=host, port=port, dbname=dbname, user=user, password=password
        )
        return conn
    except Exception as e:
        st.error(f"数据库连接失败: {e}")
        return None

def get_all_table_names(conn):
    query = "SELECT table_name FROM information_schema.tables WHERE table_schema='public';"
    with conn.cursor() as cur:
        cur.execute(query)
        result = cur.fetchall()
    return [row[0] for row in result]

def get_table_structure(db_params, table_name):
    query = """
    SELECT
        a.attname AS column_name,
        pg_catalog.format_type(a.atttypid, a.atttypmod) AS data_type,
        a.attnotnull AS is_not_null,
        a.atthasdef AS has_default,
        pg_get_expr(d.adbin, d.adrelid) AS default_value,
        pg_get_serial_sequence(c.relname, a.attname) IS NOT NULL AS is_serial,
        CASE
            WHEN tc.contype = 'p' THEN 'PRIMARY KEY'
            WHEN tc.contype = 'u' THEN 'UNIQUE'
            WHEN tc.contype = 'f' THEN 'FOREIGN KEY'
            WHEN tc.contype = 'c' THEN 'CHECK'
            ELSE NULL
        END AS constraint_type,
        fk_table.relname AS fk_table,
        fk_column.attname AS fk_column
    FROM
        pg_attribute a
    JOIN pg_class c ON a.attrelid = c.oid
    JOIN pg_namespace n ON c.relnamespace = n.oid
    LEFT JOIN pg_attrdef d ON a.attrelid = d.adrelid AND a.attnum = d.adnum
    LEFT JOIN (
        SELECT
            conrelid,
            contype,
            confrelid,
            conkey,
            confkey
        FROM
            pg_constraint
    ) tc ON c.oid = tc.conrelid AND a.attnum = ANY(tc.conkey)
    LEFT JOIN pg_class fk_table ON tc.confrelid = fk_table.oid
    LEFT JOIN pg_attribute fk_column ON fk_table.oid = fk_column.attrelid AND fk_column.attnum = ANY(tc.confkey)
    WHERE
        c.relname = %s
        AND a.attnum > 0
        AND NOT a.attisdropped
    ORDER BY
        a.attnum;
    """
    try:
        conn = psycopg2.connect(**db_params)
        with conn.cursor() as cursor:
            cursor.execute(query, (table_name,))
            columns = cursor.fetchall()
        conn.close()
        return columns
    except Exception as e:
        st.error(f"数据库查询失败: {e}")
        return None

def format_table_structure(columns):
    formatted_columns = []
    for col in columns:
        (
            column_name, data_type, is_not_null, has_default,
            default_value, is_serial, constraint_type,
            fk_table, fk_column
        ) = col

        column_def = f"    {column_name} {data_type}"

        if is_not_null:
            column_def += " NOT NULL"
        if has_default and default_value is not None:
            column_def += f" DEFAULT {default_value}"
        if constraint_type == 'PRIMARY KEY':
            column_def += " PRIMARY KEY"
        if constraint_type == 'UNIQUE':
            column_def += " UNIQUE"
        if constraint_type == 'FOREIGN KEY' and fk_table:
            ref_column = fk_column if fk_column else 'id'
            column_def += f" REFERENCES {fk_table}({ref_column})"

        formatted_columns.append(column_def)

    return "(\n" + ",\n".join(formatted_columns) + "\n)"

def get_multiple_table_structures(db_params, tables):
    structures = ""
    for table in tables:
        columns = get_table_structure(db_params, table)
        formatted = format_table_structure(columns)
        structures += f"{table} {formatted}\n\n"
    return structures

def get_completion(table_structures, sql_requirements, model="gpt-3.5-turbo"):
    instruction = "你是专业SQL工程师，请根据表结构和需求，生成高效合理的SQL查询语句。"
    prompt = f"""{instruction}

表结构：
{table_structures}

用户需求：
{sql_requirements}
"""
    messages = [{"role": "user", "content": prompt}]
    response = client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=0,
        n=1,
    )
    return response.choices[0].message.content

def download_sql(sql_text):
    output = BytesIO()
    output.write(sql_text.encode('utf-8'))
    output.seek(0)
    return output

# ------------------------------
# 页面布局
# ------------------------------

with st.sidebar:
    st.title("🔧 数据库连接")

    host = st.text_input("主机", value="localhost")
    port = st.text_input("端口", value="5432")
    dbname = st.text_input("数据库名", value="postgres")
    user = st.text_input("用户名", value="postgres")
    password = st.text_input("密码", type="password")

    if st.button("连接数据库"):
        with st.spinner("连接中..."):
            conn = connect_db(host, port, dbname, user, password)
            if conn:
                st.session_state.conn = conn
                st.session_state.db_params = {
                    "host": host,
                    "port": port,
                    "dbname": dbname,
                    "user": user,
                    "password": password
                }
                st.session_state.table_names = get_all_table_names(conn)
                st.success("✅ 数据库连接成功！")
            else:
                st.session_state.conn = None

# ------------------------------
# 主体部分
# ------------------------------

st.title("💡 PGSQL代码生成器")
st.caption("作者：何双新 ｜ 功能：PGSQL代码生成器")

if st.session_state.conn:
    st.subheader("📂 选择表并生成SQL")

    selected_tables = st.multiselect(
        "选择表（支持多选）",
        options=st.session_state.table_names,
        default=st.session_state.selected_tables
    )
    st.session_state.selected_tables = selected_tables

    if selected_tables:
        if st.button("📋 获取表结构"):
            with st.spinner("读取表结构中..."):
                st.session_state.table_structures = get_multiple_table_structures(st.session_state.db_params, selected_tables)
            st.success("✅ 表结构获取成功！")

        if st.session_state.table_structures:
            st.text_area("表结构", value=st.session_state.table_structures, height=300)

        sql_requirements = st.text_area("✍️ 输入你的需求", placeholder="例如：查询近30天新增订单数", height=150)

        if st.button("⚡ 生成SQL语句"):
            if sql_requirements and st.session_state.table_structures:
                with st.spinner("SQL生成中..."):
                    sql_result = get_completion(st.session_state.table_structures, sql_requirements)
                st.success("✅ SQL生成成功")
                st.code(sql_result, language="sql")

                st.download_button(
                    "📥 下载SQL文件",
                    data=download_sql(sql_result),
                    file_name="generated_query.sql",
                    mime="text/sql",
                )

                if sql_result.strip().lower().startswith("select"):
                    st.subheader("📊 查询结果预览")
                    try:
                        with st.session_state.conn.cursor() as cur:
                            cur.execute(sql_result)
                            df = pd.DataFrame(cur.fetchall(), columns=[desc[0] for desc in cur.description])
                        st.dataframe(df, use_container_width=True)
                    except Exception as e:
                        st.error(f"查询执行失败: {e}")
            else:
                st.warning("请输入需求描述并确保表结构已获取。")
    else:
        st.info("请选择至少一个表")
else:
    st.warning("请先连接数据库")

# ------------------------------
# 页脚
# ------------------------------
st.markdown("---")
st.caption("🚀 Streamlit + OpenAI + PostGreSQL | 最后更新：2025年4月")
st.caption("安徽智加数字科技有限公司 · 技术学习组出品 🚀")
